{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "77a83415",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import math\n",
    "from collections import Counter\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    " \n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e2b02a9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleEncoder(nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_size, dropout=0.2):\n",
    "        super(SimpleEncoder, self).__init__()\n",
    "        self.embed = nn.Embedding(vocab_size, hidden_size)\n",
    "        self.rnn = nn.GRU(hidden_size, hidden_size, batch_first=True)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x, lengths):\n",
    "        # 把batch里面的seq按照长度排序, 是pack_padded_sequence要求的\n",
    "        sorted_len, sorted_idx = lengths.sort(0, descending=True)\n",
    "        x_sorted = x[sorted_idx]\n",
    "        embedded = self.dropout(self.embed(x_sorted))\n",
    "        \n",
    "        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, sorted_len.long().cpu().data.numpy(), batch_first=True)\n",
    "        packed_out, hid = self.rnn(packed_embedded)\n",
    "        out, _ = nn.utils.rnn.pad_packed_sequence(packed_out, batch_first=True)\n",
    "        \n",
    "        _, original_idx = sorted_idx.sort(0, descending=False)\n",
    "        out = out[original_idx.long()].contiguous()\n",
    "        hid = hid[:, original_idx.long()].contiguous()\n",
    "        \n",
    "        return out, hid[[-1]]\n",
    "        \n",
    "class SimpleDecoder(nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_size, dropout=0.2):\n",
    "        super(SimpleDecoder, self).__init__()\n",
    "        self.embed = nn.Embedding(vocab_size, hidden_size)\n",
    "        self.rnn = nn.GRU(2*hidden_size, hidden_size, batch_first=True)\n",
    "        self.fc = nn.Linear(2*hidden_size, vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, y, y_lengths, hid):\n",
    "        sorted_len, sorted_idx = y_lengths.sort(0, descending=True)\n",
    "        y_sorted = y[sorted_idx]  # 隐状态也要调顺序\n",
    "        embedded = self.dropout(self.embed(y_sorted))  # batch_size, y_lengths, hidden_size\n",
    "        hid = hid[:, sorted_idx]\n",
    "\n",
    "        embedded = torch.cat([embedded, hid.squeeze(0).unsqueeze(1).expand_as(embedded)], 2)  # batch_size, y_lengths, hidden_size*2\n",
    "        \n",
    "        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, sorted_len.long().cpu().data.numpy(), batch_first=True)\n",
    "        packed_out, hid2 = self.rnn(packed_embedded, hid)\n",
    "        out, _ = nn.utils.rnn.pad_packed_sequence(packed_out, batch_first=True)\n",
    "        \n",
    "        _, original_idx = sorted_idx.sort(0, descending=False)\n",
    "        out = out[original_idx.long()].contiguous()  # batch_size, y_lengths, hidden_size\n",
    "        \n",
    "        hid2 = hid2[:, original_idx.long()].contiguous()  # 1, batch_size, hidden_size, 隐状态本来就没有长度维度\n",
    "        \n",
    "        out = torch.cat([out, hid.squeeze(0).unsqueeze(1).expand_as(out)], 2)\n",
    "        out = self.fc(out)  # batch_size, y_lengths, vocab_size\n",
    "        out = nn.functional.log_softmax(out, -1)  # log_softmax默认对第0个维度进行softmax，用-1指定为最后1维\n",
    "        return out, hid2\n",
    "    \n",
    "    \n",
    "class SimpleSeq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super(SimpleSeq2Seq, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        \n",
    "    def forward(self, x, x_lengths, y, y_lengths):\n",
    "        encoder_out, hid = self.encoder(x, x_lengths)\n",
    "        output, hid = self.decoder(y, y_lengths, hid)\n",
    "        return output, None\n",
    "    \n",
    "    def translate(self, x, x_lengths, y, max_length=10):\n",
    "        encoder_out, hid = self.encoder(x, x_lengths)\n",
    "        preds = []\n",
    "        batch_size = x.shape[0]\n",
    "        attns = []\n",
    "        for i in range(max_length):\n",
    "            output, pid = self.decoder(y=y, y_lengths=torch.ones(batch_size).long().to(device), hid=hid)\n",
    "            y = output.max(2)[1].view(batch_size, 1)\n",
    "            preds.append(y) \n",
    "        return torch.cat(preds, 1), None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "591d64d3",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'int' object has no attribute 'sort'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Input \u001b[0;32mIn [15]\u001b[0m, in \u001b[0;36m<cell line: 15>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     20\u001b[0m tto \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mLongTensor(tto)\n\u001b[1;32m     21\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[0;32m---> 22\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[43msrc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtti\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     24\u001b[0m tto \u001b[38;5;241m=\u001b[39m tto\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     25\u001b[0m output \u001b[38;5;241m=\u001b[39m output\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, output\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n",
      "Input \u001b[0;32mIn [9]\u001b[0m, in \u001b[0;36mSimpleSeq2Seq.forward\u001b[0;34m(self, x, x_lengths, y, y_lengths)\u001b[0m\n\u001b[1;32m     61\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, x_lengths, y, y_lengths):\n\u001b[0;32m---> 62\u001b[0m     encoder_out, hid \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoder\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_lengths\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     63\u001b[0m     output, hid \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdecoder(y, y_lengths, hid)\n\u001b[1;32m     64\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m output, \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniforge3/envs/tf/lib/python3.8/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1131\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "Input \u001b[0;32mIn [9]\u001b[0m, in \u001b[0;36mSimpleEncoder.forward\u001b[0;34m(self, x, lengths)\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, lengths):\n\u001b[1;32m      9\u001b[0m     \u001b[38;5;66;03m# 把batch里面的seq按照长度排序, 是pack_padded_sequence要求的\u001b[39;00m\n\u001b[0;32m---> 10\u001b[0m     sorted_len, sorted_idx \u001b[38;5;241m=\u001b[39m \u001b[43mlengths\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msort\u001b[49m(\u001b[38;5;241m0\u001b[39m, descending\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m     11\u001b[0m     x_sorted \u001b[38;5;241m=\u001b[39m x[sorted_idx]\n\u001b[1;32m     12\u001b[0m     embedded \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdropout(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed(x_sorted))\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'int' object has no attribute 'sort'"
     ]
    }
   ],
   "source": [
    "\n",
    "dropout = 0.2\n",
    "hidden_size = 128\n",
    "\n",
    "encoder = SimpleEncoder(vocab_size=len(data_c2i), hidden_size=hidden_size, dropout=dropout)\n",
    "decoder = SimpleDecoder(vocab_size=len(taget_i2c), hidden_size=hidden_size, dropout=dropout)\n",
    "model = SimpleSeq2Seq(encoder, decoder)\n",
    "model = model.to(device)\n",
    "\n",
    "# loss_fn = LanguageModelCriterion().to(device)\n",
    "# optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
    "\n",
    "for i in range(10):\n",
    "    train_iter = tain_iter()\n",
    "    for i, (src, tti, tto) in enumerate(train_iter):\n",
    "        src = torch.LongTensor(src)\n",
    "        tti = torch.LongTensor(tti)\n",
    "        tto = torch.LongTensor(tto)\n",
    "        optimizer.zero_grad()\n",
    "        output = model.forward(src, 10, tti, 10)\n",
    "\n",
    "        tto = tto.reshape(-1)\n",
    "        output = output.reshape(-1, output.shape[-1])\n",
    "        loss = criterion(output, tto)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(f\"{loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b76fd7c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79df57d1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba6ae87d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "067015b4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5595c3e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d7ffaac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "11780812",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "data_path = \"letters_source.txt\"\n",
    "target_path = \"letters_target.txt\"\n",
    "with open(data_path, 'r') as f:\n",
    "    data_s = f.read().split('\\n')\n",
    "    \n",
    "with open(target_path, 'r') as f:\n",
    "    target_s = f.read().split('\\n')\n",
    "    \n",
    "data_chars = sorted(list(set(\"\".join(data_s))))\n",
    "taget_chars = sorted(list(set(\"\".join(target_s))))\n",
    "pad_flage = '<PAD>'\n",
    "unk_flage = '<UNK>'\n",
    "beg_flage = '<GO>'\n",
    "end_flage = '<EOS>'\n",
    "data_chars = [pad_flage, unk_flage, beg_flage, end_flage] + data_chars\n",
    "taget_chars = [pad_flage, unk_flage, beg_flage, end_flage] + taget_chars\n",
    "\n",
    "data_c2i = {char: i for i, char in enumerate(data_chars)}\n",
    "taget_c2i = {char: i for i, char in enumerate(taget_chars)}\n",
    "taget_i2c = {i: char for i, char in enumerate(taget_chars)}\n",
    "\n",
    "max_len = 10\n",
    "\n",
    "data = []\n",
    "for line_ in data_s:\n",
    "    line = [data_c2i.get(char, 1) for char in  line_]\n",
    "    line = line[:max_len]\n",
    "    line = line + [0]*(max_len - len(line))\n",
    "    data.append(line)\n",
    "target_i = []\n",
    "target_o = []\n",
    "for line_ in target_s:\n",
    "    line = [2] + [taget_c2i.get(char, 1) for char in  line_] + [3]\n",
    "    line = line[:max_len]\n",
    "    line = line + [0]*(max_len - len(line))\n",
    "    \n",
    "    target_i.append(line)\n",
    "    line = [taget_c2i.get(char, 1) for char in  line_] + [3]\n",
    "    line = line[:max_len]\n",
    "    line = line + [0]*(max_len - len(line))\n",
    "    target_o.append(line)\n",
    "    \n",
    "    \n",
    "data = np.array(data)\n",
    "target_i = np.array(target_i)\n",
    "target_o = np.array(target_o)\n",
    "\n",
    "split_index = int(len(data)*0.7)\n",
    "train_data = data[:split_index]\n",
    "train_target_i = target_i[:split_index]\n",
    "train_target_o = target_o[:split_index]\n",
    "\n",
    "test_data = data[split_index:]\n",
    "target_data = target_o[split_index:]\n",
    "\n",
    "def tain_iter(batch_size=128):\n",
    "    num = len(train_data)//batch_size\n",
    "    for i in range(num):\n",
    "        beg_index = i*batch_size\n",
    "        end_index = (i+1)*batch_size\n",
    "        td = train_data[beg_index: end_index]\n",
    "        tti = train_target_i[beg_index: end_index]\n",
    "        tto = train_target_o[beg_index: end_index]\n",
    "        yield td, tti, tto\n",
    "\n",
    "def test_iter(batch_size=128):\n",
    "    num = len(test_data)//batch_size\n",
    "    for i in range(num):\n",
    "        beg_index = i*batch_size\n",
    "        end_index = (i+1)*batch_size\n",
    "\n",
    "        yield test_data[beg_index: end_index], target_data[beg_index: end_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2388f5cf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
